Skip to content

Centralizes environment variable access by routing variables reads through the envs.py module #1147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 8 commits into
base: main
Choose a base branch
from
Open

Centralizes environment variable access by routing variables reads through the envs.py module #1147

wants to merge 8 commits into from

Conversation

@xingliu14
Copy link
Collaborator

@xingliu14 xingliu14 commented Nov 21, 2025 •
edited
Loading

Description

Centralizes environment variable access by routing following environment variable reads through the envs.py module.

  • SKIP_JAX_PRECOMPILE
  • VLLM_XLA_CHECK_RECOMPILATION
  • MODEL_IMPL_TYPE
  • NEW_MODEL_DESIGN
  • PHASED_PROFILING_DIR
  • PYTHON_TRACER_LEVEL
  • USE_MOE_EP_KERNEL
  • NUM_SLICES
  • RAY_USAGE_STATS_ENABLED
  • VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE

Addressing #1016.

Tests

pytest

tests/runner/test_utils.py
tests/worker/tpu_worker_test.py
tests/runner/test_tpu_runner.py
tests/layers/common/test_attention_interface.py
tests/executors/test_ray_distributed_executor.py
tests/runner/test_tpu_runner_mesh.py
tests/test_envs.py

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@xingliu14 xingliu14 marked this pull request as draft November 21, 2025 05:33
@xingliu14 xingliu14 force-pushed the env_var branch 4 times, most recently from 4df238b to 2c92d5d Compare November 21, 2025 07:04
@xingliu14 xingliu14 marked this pull request as ready for review November 21, 2025 07:08
@xingliu14
Copy link
Collaborator Author

xingliu14 commented Nov 21, 2025

@kyuyeunk It's ready for review.

kyuyeunk
kyuyeunk reviewed Nov 22, 2025
View reviewed changes
examples/offline_inference.py Outdated
if __name__ == "__main__":
# Skip long warmup for local simple test.
os.environ['SKIP_JAX_PRECOMPILE'] = '1'
envs.environment_variables['SKIP_JAX_PRECOMPILE'] = lambda: True
Copy link
Collaborator

@kyuyeunk kyuyeunk Nov 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this seems unintuitive way to specify environment variable within unit tests.

does something like envs.SKIP_JAX_PRECOMPILE = True not work?

Copy link
Collaborator Author

@xingliu14 xingliu14 Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is no setattr to write the variable. Actually the best to rewrite is setting the environ direcly. I have reverted the changes along with other similar changes.

tpu_inference/executors/ray_distributed_executor.py Outdated
from vllm.v1.executor.ray_executor import RayWorkerMetaData
from vllm.v1.executor.ray_utils import RayWorkerWrapper, _wait_until_pg_ready

import tpu_inference.envs as tpu_envs
Copy link
Collaborator

@kyuyeunk kyuyeunk Nov 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

other files just use envs for tpu_inference.envs and reword vllm.envs as vllm_envs. what is the rationale for doing it differently here?

Copy link
Collaborator Author

@xingliu14 xingliu14 Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a mistake, and actually I reverted the file change since it's var setting instead of reading.

@xingliu14
Update offline_lora_inference.py to use tpu_inference.envs
d4a3ea0 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
Refactor environment variable handling to use EnvironmentVariables class
0e17a51 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
fix
7679f62 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
Revert test file changes
f5b1ddc 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
Fix isort, yapf, and ruff formatting issues
8d13d94 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
Use os.environ directly for setting environment variables
3e8e3db 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
fix
2646f20 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
@xingliu14
revert tpu_worker_test.py
273ea3d 
Signed-off-by: Xing Liu <xingliu14@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kyuyeunk kyuyeunk kyuyeunk left review comments

@sixiang-google sixiang-google Awaiting requested review from sixiang-google sixiang-google is a code owner

@mrjunwan-lang mrjunwan-lang Awaiting requested review from mrjunwan-lang mrjunwan-lang is a code owner

@py4 py4 Awaiting requested review from py4 py4 is a code owner

@vanbasten23 vanbasten23 Awaiting requested review from vanbasten23 vanbasten23 is a code owner

@wenxindongwork wenxindongwork Awaiting requested review from wenxindongwork wenxindongwork is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@xingliu14 @kyuyeunk